#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copyright 2020 Huawei Technologies Co., Ltd
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import os
import sys
import time
import copy
import numpy as np
from queue import Queue
import acl
import live555
from .dvpp import DvppInfo
sys.path.append("..")
from common import const
from common.log import Log,check_ret
from resource.resource import ascend_resource


class AscendStreamPuller(object):
    def __init__(self, channelID, streamName, video_path, q, width, height, context):
        self.class_name = self.__class__.__name__
        # live555 info
        self.handle = None
        self.channelID = channelID
        self.streamName = streamName
        # video path(now it is not usefull)
        self.video_path = video_path
        # queue for communication between threads
        self.queue = q 
        # frame info
        self.context = context
        self.channelID = channelID
        self.width = width
        self.height = height
        self.frame_id = 0
        self.device_ptr = 0
        self.size = 0
        # frame payload
        self.frame_dict = {'sps': None, 'pps': None, 'sei': None, 'i': None, 'p': None}
        # ctrl
        self.stop_lock = False
        self.can_stop = False
        # memory
        self.memory = ascend_resource.memory_pool

    def start_stream(self):
        useTCP = True
        self.handle = live555.startRTSP(self.streamName, 
                                 self.oneFrame, 
                                 self.shutdownCallback, 
                                 useTCP)
        if self.handle is None:
            check_ret(self.class_name, 'start_stream',\
                'fail to start rstsp(channelID{})'.format(self.channelID), self.handle)
        
    def stop_stream(self):
        self.stop_lock = True
        time.sleep(0.2)
        while True:
            if self.can_stop:
                break
        try:
            live555.stopRTSP(self.handle)
        except live555.error as e:
            Log.do_log(WARNING, 'StreamPuller', 'stop_stream', 'live555 error received')

    def shutdownCallback(self):
        print('shutdown callback')
  
    def transfer_host_to_device(self, frame_data):
        ret = acl.rt.set_context(self.context)
        check_ret(self.class_name, 'transfer_host_to_device',\
            'fail to set context', ret)
        ptr = acl.util.numpy_to_ptr(frame_data)
        self.size = frame_data.size * frame_data.itemsize
        self.device_ptr = self.memory.memory_allocation_(self.size, const.DVPP)
        ret = acl.rt.memcpy(self.device_ptr, 
                            self.size, 
                            ptr, 
                            self.size, 
                            const.ACL_MEMCPY_HOST_TO_DEVICE)
        check_ret(self.class_name, 'transfer_host_to_device',\
            'fail to memcpy from host to device', ret)

    def send_frame(self, frame_data):
        self.transfer_host_to_device(frame_data)
        video_data = DvppInfo()
        video_data.buffer_ptr = self.device_ptr
        video_data.buffer_size = self.size
        video_data.width = self.width
        video_data.height = self.height
        self.frame_id += 1
        video_data.frame_id = self.frame_id
        video_data.channel_id = self.channelID
        self.queue.put(video_data)
        self.device_ptr = 0

    def set_one_frame(self, name, data):
        if name == 'sei' or name == 'i':
            header = np.array([0, 0, 1], np.uint8)
        else:
            header = np.array([0, 0, 0, 1], np.uint8)
        one_frame = np.concatenate((header, data))
        self.frame_dict[name] = one_frame

    def set_frame(self, np_data):
        if np_data[0] == 103:
            self.clear_frame()
            self.set_one_frame('sps', np_data)
            ctrl = 1
        elif np_data[0] == 104:
            self.set_one_frame('pps', np_data)
            ctrl = 2
        elif np_data[0] == 6:
            self.set_one_frame('sei', np_data)
            ctrl = 4
        elif np_data[0] == 101:
            self.set_one_frame('i', np_data)
            ctrl = 8
        elif np_data[0] == 65:
            self.set_one_frame('p', np_data)
            ctrl = 16
        else:
            ctrl = 0
        return ctrl  

    def clear_frame(self):
        for name in self.frame_dict.keys():
            self.frame_dict[name] = None
        self.ctrl = 0

    def concat_frame(self, ctrl):
        self.ctrl |= ctrl
        if self.ctrl == 15:
            group = np.concatenate((self.frame_dict['sps'], \
                                    self.frame_dict['pps'], \
                                    self.frame_dict['sei'], \
                                    self.frame_dict['i']))
        elif self.ctrl == 11:
            group = np.concatenate((self.frame_dict['sps'], \
                                    self.frame_dict['pps'], \
                                    self.frame_dict['i']))
        elif self.ctrl == 27 or self.ctrl == 31:
            group = self.frame_dict['p']
        else:
           group = None
        return group

    def oneFrame(self, codecName, data, sec, usec, durUSec):
        if self.stop_lock:
            return
        if self.queue is None or \
           self.queue.full():
            self.clear_frame()
            return
        if codecName != 'H264':
            return
        # read frame data
        np_data = np.frombuffer(data, dtype=np.uint8)
        # set frame data
        ctrl = self.set_frame(np_data)    
        # concatenate frame data               
        frame = self.concat_frame(ctrl) 
        #print('frame: {}, ctrl: {}, self.ctrl: {}'.format(frame, ctrl, self.ctrl))
        # send data
        if frame is not None:
            self.can_stop = False
            self.send_frame(frame)
            self.can_stop = True


def run_stream():
    live555.runEventLoop()

def del_stream():
    live555.stopEventLoop()
